from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from langchain_core.output_parsers import JsonOutputParser


class SubScheme(BaseModel):
    # project_name: str = Field(description="项目的名称")
    # abstract: str = Field(description="全文的摘要")
    # title: str = Field(description="一级标题")
    # sub_title: str = Field(description="二级标题")
    content: str = Field(description="生成的内容")


_writer_system_template = """
作为一名杰出的方案撰写专家，您的任务是将项目的核心价值和愿景通过精心编排的内容展现给读者。您需要依据项目名称、全文摘要以及精心设计的一级和二级标题，来构思并撰写目录下的内容。以下是对您职责的精炼描述：

您的工作是将项目的精髓和目标通过以下结构化方式传达给目标受众：

1. **项目名称**：这是您内容构建的起点，一个响亮且富有内涵的项目名称能够立即吸引读者的注意。

2. **全文摘要**：在这一部分，您需要简洁明了地概述项目的核心内容和主要目标，为读者提供一个关于项目全貌的预览。

3. **一级标题**：这些是您内容的主要章节，每个标题都应代表一个关键的主题或项目的关键组成部分。

4. **二级标题**：在每个一级标题下，您将细化为更具体的小节，每个二级标题都深入探讨一级标题下的具体内容或子主题。

在撰写过程中，您需要确保每个部分都紧密相连，共同构建一个有说服力和连贯性的叙述。您的文字应该清晰、准确，同时富有吸引力，能够激发读者的兴趣，并引导他们深入了解项目的细节。通过您的专业写作技巧，您将把枯燥的信息转化为引人入胜的故事，让每个读者都能理解和感受到项目的价值和潜力。

要求:
作为一名杰出的方案撰写专家，您的角色不仅仅是文字的搬运工，更是故事的编织者和愿景的传达者。以下是对您职责的进一步细化和要求：

1. **项目名称**：
   - 确保项目名称具有独特性，能够准确反映项目的本质和目标市场。
   - 名称应简洁、易于记忆，同时具有较强的品牌潜力。

2. **全文摘要**：
   - 提供一个精炼的概述，突出项目的关键卖点和创新之处。
   - 摘要应足够吸引人，激发读者进一步阅读全文的兴趣。

3. **一级标题**：
   - 每个一级标题都应代表一个清晰的主要部分，涵盖项目的一个关键方面或阶段。
   - 标题应具有逻辑性和层次感，引导读者顺畅地理解整个方案的结构。

4. **二级标题**：
   - 在每个一级标题下，二级标题应深入探讨具体的子主题或细节，提供更详尽的信息。
   - 二级标题应具体、明确，使读者能够快速定位感兴趣的内容。

5. **内容撰写要求**：
   - 确保内容的准确性和真实性，所有数据和信息都应经过核实。
   - 语言应清晰、简洁，避免行业术语或复杂词汇，确保非专业读者也能理解。
   - 内容应具有逻辑性和连贯性，每个部分都应紧密相连，形成一个统一的整体。
   - 通过故事讲述的方式，使内容更具吸引力，让读者能够与项目产生情感共鸣。
   - 在适当的地方使用图表、图像和其他视觉元素来辅助说明，增强内容的吸引力和易读性。
   - 确保方案的格式和风格符合专业标准，包括正确的语法、拼写和标点。
   - 在撰写过程中，保持对项目目标和预期结果的清晰认识，确保所有内容都围绕这些目标展开。

6. **审校和反馈**：
   - 在完成初稿后，进行彻底的自我审校，确保内容无误。
   - 愿意接受同事、客户或目标读者的反馈，并据此改进内容。
   - 在最终提交前，确保方案经过多轮校对和修改，以达到最高质量标准。

7. 不要有序号

8.输出的内容不要带*号


通过遵循这些要求，您将能够创作出既专业又具有吸引力的方案，有效地传达项目的价值和愿景，同时激发读者的兴趣和参与。


输出：
{output_format}
"""

_writer_human_template = """
项目标题：{project_name}
全文摘要:{abstract}
一级目录:{title}
二级目录:{sub_title}
"""


class Writer:

    def __init__(self, llm):

        self._llm = llm

        _prompt = ChatPromptTemplate.from_messages(
            [("system", _writer_system_template), ("human", _writer_human_template)]
        )

        _parser = JsonOutputParser(pydantic_object=SubScheme)
        _prompt = _prompt.partial(output_format=_parser.get_format_instructions())

        self._chain = _prompt | self._llm | _parser

    def __call__(self, state):
        return self._chain.invoke(state)


if __name__ == "__main__":

    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        base_url="http://192.168.10.11:60026/v1",
        model="qwen2.5:7b",
        api_key="ollama",
    )

    write = Writer(_llm)
    rt = write(
        {
            "project_name": "基于多模态大模型驱动的数字人",
            "abstract": "本方案旨在设计和开发一个基于多模态大模型的数字人系统，结合文本、语音、视觉等多模态数据提供智能交互能力。通过先进的训练技术和优化方法，该数字能够在各种场景中展现丰富的人类行为、情感和语言交流特质，实现高度智能化且自然互动。",
            "sub_title": "团队结构与角色划分",
            "title": "体实施和开发流程",
        }
    )
    print(rt["content"])
